from data_preprocess import *
from train_model import *
from causalnex.discretiser import Discretiser

def to_discret(data):
    data_discret = data.copy()
    data_vals = {col: data[col].unique() for col in data.columns}

    failures_map = {v: 'no-failure' if v == [0] else 'have-failure' for v in data_vals['failures']}
    studytime_map = {v: 'short-studytime' if v in [1, 2] else 'long-studytime' for v in data_vals['studytime']}
    data_discret["failures"] = data_discret["failures"].map(failures_map)
    data_discret["studytime"] = data_discret["studytime"].map(studytime_map)

    data_discret["absences"] = Discretiser(method="fixed",
                                               numeric_split_points=[1, 10]).transform(
        data_discret["absences"].values)

    data_discret["G1"] = Discretiser(method="fixed",
                                         numeric_split_points=[10]).transform(data_discret["G1"].values)

    data_discret["G2"] = Discretiser(method="fixed",
                                         numeric_split_points=[10]).transform(data_discret["G2"].values)

    data_discret["G3"] = Discretiser(method="fixed",
                                         numeric_split_points=[10]).transform(data_discret["G3"].values)
    absences_map = {0: "No-absence", 1: "Low-absence", 2: "High-absence"}

    G1_map = {0: "Fail", 1: "Pass"}
    G2_map = {0: "Fail", 1: "Pass"}
    G3_map = {0: "Fail", 1: "Pass"}

    data_discret["absences"] = data_discret["absences"].map(absences_map)
    data_discret["G1"] = data_discret["G1"].map(G1_map)
    data_discret["G2"] = data_discret["G2"].map(G2_map)
    data_discret["G3"] = data_discret["G3"].map(G3_map)

    return data_discret


if __name__ == '__main__':
    data = get_csv_data('student-por.csv',['school','sex','age','Mjob', 'Fjob','reason','guardian'])
    data = data.head(5)
    data_val = to_val(data)
    data_discret = to_discret(data)

    net = train_network(data_val)
    bn = train_bn(data_discret,net)
